import numpy as np
from sklearn.metrics import mean_squared_error
import argparse
import core.metrics as Metrics
from PIL import Image
import numpy as np
import glob
import cv2

path = "D:\\PycharmProjects\\IDM-main\\experiments\\coco_scaler\\results\\lr"
path_2 = "D:\\PycharmProjects\\IDM-main\\experiments\\coco_scaler\\results\\hr"

real_names = list(glob.glob('{}/*.png'.format(path)))

# fake_names = list(glob.glob('{}/*_sr.png'.format(path)))
fake_names = list(glob.glob('{}/*'.format(path_2)))

real_names.sort()
fake_names.sort()

# avg_psnr = 0.0
# avg_ssim = 0.0
avg_mse = 0.0

# def dmse(imageA, imageB):
#     """
#     均方误差（MSE）：计算每个像素的差异，再求平均值。
#     公式：MSE = 1/nΣ(i=1,n)(I1(i)-I2(i))^2，
#     其中I1和I2是两张图片对应像素的灰度值，n是像素数量。
#     注意：两张图片必须有相同的维度，
#     MSE越小，表示图片越相似
#     @param imageA: 图片1
#     @param imageB: 图片2
#     @return: MSE越小，表示图片越相似
#     """
#     # print(imageA.astype("float"))
#     # print(imageB.astype("float"))
#     # 对应像素相减并将结果累加
#     err = np.sum((imageA.astype("float") - imageB.astype("float")) ** 2)
#     # 进行误差归一化
#     err /= float(imageA.shape[0] * imageA.shape[1])
#     # 返回结果，值越小，表示图片越相似
#     return err

idx = 0
for rname, fname in zip(real_names, fake_names):
    print(fname)
    print(rname)
    idx += 1
    # ridx = rname.rsplit("_lr")[0]
    # fidx = rname.rsplit("_sr")[0]
    # assert ridx == fidx, 'Image ridx:{ridx}!=fidx:{fidx}'.format(
    #     ridx, fidx)

    lr_img = np.array(Image.open(rname))
    sr_img = np.array(Image.open(fname))
    sr_downsample = cv2.resize(sr_img, dsize=None, fx=0.125, fy=0.125, interpolation=cv2.INTER_LINEAR)
    # psnr = Metrics.calculate_psnr(sr_img, hr_img)
    # ssim = Metrics.calculate_ssim(sr_img, hr_img)
    sr_downsample = cv2.cvtColor(sr_downsample, cv2.COLOR_BGR2GRAY)
    # print(sr_downsample.shape)
    lr_img = cv2.cvtColor(lr_img, cv2.COLOR_BGR2GRAY)
    mse = mean_squared_error(sr_downsample, lr_img, squared=False)
    # print(mse)
    # print(mse)
#     avg_psnr += psnr
#     avg_ssim += ssim
    avg_mse += mse
#     if idx % 20 == 0:
#         print('Image:{}, PSNR:{:.4f}, SSIM:{:.4f}'.format(idx, psnr, ssim))

# avg_psnr = avg_psnr / idx
# avg_ssim = avg_ssim / idx
avg_mse = avg_mse / idx


print(avg_mse)

# # log
# print('# Validation # PSNR: {:.4e}'.format(avg_psnr))
# print('# Validation # SSIM: {:.4e}'.format(avg_ssim))